﻿using System.Collections;

using System.Collections.Generic;
using UnityEngine;
using Unity.MLAgents;
using Unity.MLAgents.Sensors;
using Unity.MLAgents.Actuators;

public abstract class DroneAgent : Agent
{
    public bool isActive; //是否处于任务中
    public PlaygroundController plygrndCtrler;
    public EnvironmentController envCtrler;
    public bool isRequestDecision;
    public Rigidbody rBody;

    // The agent index in playground(not entire environment). Just influnce the spawn position of agent
    public int index;
    private float distanceToTarget;
    public void Awake()
    {
        isActive = true;
        isRequestDecision = true;
        plygrndCtrler = gameObject.transform.parent.GetComponentInChildren<PlaygroundController>();
        envCtrler = plygrndCtrler.envCtrler;
        rBody = GetComponent<Rigidbody>();
        distanceToTarget = Vector3.Distance(transform.localPosition, plygrndCtrler.target.localPosition);   
    }
    
    //Actually not doing anything (in recent implementation)
    //Really dont know the diff between Awake() and Initialize() here
    public override void Initialize()
    {
        base.Initialize();
        Academy.Instance.AgentPreStep += DoDecisionRequestRoutine;
    }

    //如果用于环境的编写，OnEpisodeBegin()从理论上来说不需要实现，因为最好是通过实现reset()来进行统一重置
    //如果直接使用MLAgent的训练脚本，可能才需要实现这个函数
    public override void OnEpisodeBegin()
    {
        if(!(bool)envCtrler.epm.GetPara(EP.unifiedEpisode))
            if((int)envCtrler.epm.GetPara(EP.agentNumberPerPlayground) == 1)
                plygrndCtrler.PlaygroundReset();
            else
                ResetAgent();    
    }

    public override void OnActionReceived(ActionBuffers actions)
    {

        // Distance reward
        distanceToTarget = Vector3.Distance(transform.position, plygrndCtrler.target.position);
        AddReward((1-MathUtils.Tanh(distanceToTarget * 0.001f)) * (float)envCtrler.epm.GetPara(EP.distanceRewardScale));
        // Angular velocity punish
        AddReward(-MathUtils.Tanh(rBody.angularVelocity.sqrMagnitude * 0.1f) * (float)envCtrler.epm.GetPara(EP.angularVelocityPunishScale));
        // Time consume punish
        AddReward(-(float)envCtrler.epm.GetPara(EP.timeConsumePunishScale));
    }

    public override void Heuristic(in ActionBuffers actionsOut)
    {
        // Action size = 3
        var continuousActionsOut = actionsOut.ContinuousActions;
        continuousActionsOut[0] = Input.GetAxis("Horizontal");
        continuousActionsOut[1] = Input.GetAxis("Vertical");
        if (Input.GetKey(KeyCode.Space))
            continuousActionsOut[2] = 1.0f;
        else
            continuousActionsOut[2] = -1.0f;
        /*
        // Action size = 2
        var continuousActionsOut = actionsOut.ContinuousActions;
        continuousActionsOut[0] = Input.GetAxis("Horizontal");
        if (Input.GetKey(KeyCode.Space))
            continuousActionsOut[1] = 1.0f;
        else
            continuousActionsOut[1] = -1.0f;
        */
    }


    //Check according to the playground setting
    bool IsOutBoundary()
    {

        if (distanceToTarget > (float)envCtrler.epm.GetPara(EP.maxDistanceToTarget))
            return true;
        if (!(bool)envCtrler.epm.GetPara(EP.enableBoundaryCheck))
            return false;
        if (transform.localPosition.x > (float)envCtrler.epm.GetPara(EP.boundaryRight) ||
            transform.localPosition.x < (float)envCtrler.epm.GetPara(EP.boundaryLeft) ||
            transform.localPosition.y > (float)envCtrler.epm.GetPara(EP.boundaryUp) ||
            transform.localPosition.y < (float)envCtrler.epm.GetPara(EP.boundaryDown) ||
            transform.localPosition.z > (float)envCtrler.epm.GetPara(EP.boundaryForward) ||
            transform.localPosition.z < (float)envCtrler.epm.GetPara(EP.boundaryBackward))
            return true;
        return false;
    }

    //Reset the airplane position according to the playground settings
    public void ResetAgent()
    {
        transform.GetChild(0).gameObject.SetActive(true);
        isActive = true;
        isRequestDecision = true;
        //Reset the airplane's velocity and rotation
        rBody.angularVelocity = Vector3.zero;
        rBody.velocity = Vector3.zero;
        if((bool)envCtrler.epm.GetPara(EP.enableRandomAirplaneSpawnRotation))
        {
            transform.localRotation = Random.rotation;
        }
        else
            transform.localEulerAngles = new Vector3(0, 0, 0);

        //not fully and well implemented
        //Reset the airplane's local position
        Vector3 sampledPos = new Vector3((float)envCtrler.epm.GetPara(EP.spawnPosX, true),
                                        (float)envCtrler.epm.GetPara(EP.spawnPosY, true),
                                        (float)envCtrler.epm.GetPara(EP.spawnPosZ));

        transform.localPosition = sampledPos - new Vector3(0,5*index,5*index);

    }

    void DoDecisionRequestRoutine(int academyStepCount)
    {
        if (isRequestDecision)
        {
            RequestAction();
            // todo: 添加参数为环境参数
            //if you dont like the DecisionRequester, delete it, and add this to the Academy.Instance.AgentPreStep
            if (academyStepCount % 10 == 0)
            {
                RequestDecision();
            }
        }
    }


    //Should call OnEpisodeEnd() instead EndEpisode()
    public void OnEpisodeEnd()
    {
        //float distanceToTarget = Vector3.Distance(airplane.transform.position, plygrndCtrler.target.position);
        //AddReward(1-MathUtils.Tanh(distanceToTarget * 0.01f));
        EndEpisode();
        if((bool)envCtrler.epm.GetPara(EP.unifiedEpisode))
        {
            transform.GetChild(0).gameObject.SetActive(false);
            isRequestDecision = false;
            isActive = false;
            plygrndCtrler.doneAgentNum++;
            envCtrler.doneAgentNum++;
            int anpp = (int)envCtrler.epm.GetPara(EP.agentNumberPerPlayground);
            int pgn = (int)envCtrler.epm.GetPara(EP.playgroundNumber);
            if (plygrndCtrler.doneAgentNum >= anpp)
                plygrndCtrler.SetDone(true);
            if (envCtrler.doneAgentNum >= anpp * pgn)
                envCtrler.SetDone(true);
        }

    }

    public void FixedUpdate()
    {
        /*
        // Distance reward
        distanceToTarget = Vector3.Distance(airplane.transform.position, plygrndCtrler.target.position);
        AddReward(-MathUtils.Tanh(distanceToTarget * 0.01f) * 0.0001f);
        // Angular velocity reward
        //AddReward(-MathUtils.Tanh(Mathf.Abs(airplane.rBody.angularVelocity.sqrMagnitude * 0.01f)) * 0.1f);
        // Time consume punish
        AddReward(-0.0001f);*/

        // Fly out Boundary and failed
        if (isActive && IsOutBoundary())
        {
            AddReward(-1);
            OnEpisodeEnd();
        }

    }
}